import os
from PIL import Image
import torch.nn as nn
from torchvision import transforms

# from torch.utils.data import Dataset
# import numpy as np

# import warnings
# import matplotlib.pyplot as plt
# warnings.filterwarnings("ignore")
# plt.ion()

# Resize(84) 将图像较短边调整为84像素，较长边将按比例缩放
# Resize([256, 256]) 将图像宽和高分别调整为256像素
# CenterCrop(84)：从中心裁剪图片宽高为84的正方形
# ToTensor()：读取图片像素且转化为0-1的数字（进行归一化）
data_transform = {
    'train': transforms.Compose([
        transforms.Resize(84),
        transforms.CenterCrop(84),
        # 转换成tensor向量
        transforms.ToTensor(),
        # 对图像进行归一化操作
        # [0.485, 0.456, 0.406]，RGB通道的均值与标准差
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.Resize(84),
        transforms.CenterCrop(84),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}


def pil_open_img(path):
    dir_img = r'C:/Users/wbl/Desktop/pythonProject1/image/'
    img_path = os.path.join(dir_img, path)
    # 以RGB格式打开图像；Pytorch DataLoader就是使用PIL读取的图像格式
    return Image.open(img_path).convert('RGB')


class DataSetCust(nn.Module):
    def __init__(self, transform=None, loader=None):
        super(DataSetCust, self).__init__()
        self.images = []
        self.labels = []
        self.loader = loader
        self.transform = transform

    def load_txt(self, filename):
        fp = open(filename, 'r')
        self.images = []
        self.labels = []
        for line in fp:
            line.strip('\n')
            line.rstrip()
            str_arr = line.split()
            self.images.append(str_arr[0])
            self.labels.append(int(str_arr[1]))

    # 重写函数,读取数据
    def __getitem__(self, index):
        img_name = self.images[index]
        label = self.labels[index]
        image = self.loader(img_name)
        if self.transform is not None:
            image = self.transform(image)
        return image, label

    # 重写函数，返回数据集数据个数
    def __len__(self):
        return len(self.images)


def get_dataset(key: str, file: str):
    # 生成Pytorch所需的DataLoader数据输入格式
    ds = DataSetCust(transform=data_transform[key], loader=pil_open_img)
    ds.load_txt(file)
    return ds


# 验证是否生成DataLoader格式数据
def echo_dataloader(data_loader):
    for data in data_loader:
        inputs, labels = data
        print(inputs)
        print(labels)
